package com.easy.util;

import java.util.*;
import java.util.concurrent.*;

public class ExpiringMap<K, V> {
    private final Map<K, V> map = new HashMap<>();
    private final Map<K, Long> expiryTimes = new HashMap<>();
    private final ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
    private final long defaultExpiryMillis;

    public ExpiringMap(long defaultExpiryMillis) {
        this.defaultExpiryMillis = defaultExpiryMillis;
        executor.scheduleAtFixedRate(this::cleanup, 1, 1, TimeUnit.MILLISECONDS);
    }

    public void put(K key, V value) {
        put(key, value, defaultExpiryMillis);
    }

    public void put(K key, V value, long expiryMillis) {
        map.put(key, value);
        expiryTimes.put(key, System.currentTimeMillis() + expiryMillis);
    }

    public V get(K key) {
        if (!expiryTimes.containsKey(key) || expiryTimes.get(key) < System.currentTimeMillis()) {
            remove(key);
            return null;
        }
        return map.get(key);
    }
    public boolean containsKey(K key) {
        System.out.println(map);
        return map.containsKey(key);
    }

    public void remove(K key) {
        map.remove(key);
        expiryTimes.remove(key);
    }

    private void cleanup() {
        long now = System.currentTimeMillis();
        expiryTimes.entrySet().removeIf(entry -> {
            if (entry.getValue() < now) {
                map.remove(entry.getKey());
                return true;
            }
            return false;
        });
    }

    public void shutdown() {
        executor.shutdown();
    }
}
